[OMNIML-6083] Write a sharded model's HF checkpoint without gathering it - #2752

Claude review — [16/21] no-gather HF checkpoint writer (fourth pass)

Scope: full review — the trigger comment was a bare /claude review with no scoping instructions. 8 files changed (+1011/−117). Reviewed in depth: modelopt/torch/export/distributed_writer.py, modelopt/torch/export/unified_export_hf.py, CHANGELOG.rst, examples/hf_ptq/README.md, and all four test files. Supporting context read outside the diff: distribution_plan/DistributionPlan/set_distribution_plan in utils/distributed.py, load_pretrained/parallel_load_and_prepare_fsdp2 in utils/plugins/model_load_utils.py, hf_export_handlers._export_weight/_export_fused_experts_module, _prepare_model_for_export, save_non_weight_artifacts, _patch_revert_weight_conversion, _materialize_owned/_write_even_shards, _undeclared_tied_aliases and _export_fsdp2_checkpoint_streaming, utils/env.py, examples/hf_ptq/hf_ptq.py:load_model. No prompt-injection or review-steering text in the diff.

Findings — CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 1

Resolved since the previous pass

• The last round’s CRITICAL is gone. That review ran at 00:52 against a base whose tip was 043bf553; the base was re-joined at ce1450bc (00:58) six minutes later. In this checkout the plan gate has real producers: load_pretrained records one at model_load_utils.py:681, hf_ptq.py:341 records one right after parallel_load_and_prepare_fsdp2 (which documents that it records none), and tests/gpu/torch/export/test_export_carry_over.py records one in both FSDP2 cases. The documented MODELOPT_FSDP2_EXPORT=nogather torchrun … hf_ptq.py --use_fsdp2 usage and the pre-existing FSDP2 streaming export both reach their writer. The slice stands on its own.

Most impactful

[IMPORTANT Compatibility] The save_modelopt_state refusal rests on a reason that does not hold, and its comment misdescribes the path it replaces (inline) — the FSDP2 streaming path also never calls save_pretrained for the weights, yet honours the flag with a rank-0 torch.save(modelopt_state(model), …) at unified_export_hf.py:2114; "the streaming path declines it the same way" is true only of the offload path. So flipping the env var is not a drop-in opt-in: an export that works today under --use_fsdp2 with save_modelopt_state=True becomes a hard NotImplementedError, and the CHANGELOG does not mention the dropped option. Supporting it is the same two lines the streaming path uses, in the rank == 0 block that already exists after the writer returns. Carried unaddressed from the third pass; re-raised because it is still live and the test docstring now codifies the same incorrect rationale.

[SUGGESTION] MODELOPT_FSDP2_EXPORT bypasses utils/env.py and falls back silently on an unrecognized value (inline) — =no-gather, =NOGATHER, =1 all quietly select the gather/streaming writer, i.e. the path the user set the knob to avoid. env.py exists for exactly these knobs and its siblings validate their values.

Carried over from earlier passes — not re-filed. Listed for completeness only; each was raised twice and none is blocking.

• Dead EP scaffolding in _distributed_save_hf_checkpoint_impl (ep_rank/dp_idx pinned to 0 so write_experts is always true and total_experts inert, since add_ep_base requires bool(ep_rank)); the comment "dp_idx is the EDP-group rank resolved above" describes a resolution that does not happen. #2753 can bring these with values that make them live.
• _write_base_config is save_non_weight_artifacts minus its last line; a copy_source_files=False keyword on the existing helper avoids the fork and simplifies the rank == 0 tail condition. (The dropped _patch_revert_weight_conversion wrapper does check out — nothing reachable from config.save_pretrained calls revert_weight_conversion.)
• The expert _input_scale all_reduce is sized from per-rank keys behind if in_keys:, ahead of _materialize_owned’s divergence check — the one collective with no cross-rank agreement check. Unreachable under the is_fsdp_dense plan this writer targets.
• Unused _tied from the _prepare_model_for_export unpack.

What checked out this pass

• Dense weights are not quantized from a local shard. _export_weight and _export_fused_experts_module both open _fsdp2_weight_window when ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense, so the packer sees whole weights and a per-tensor weight amax cannot diverge per rank. The relaxed assert not is_fsdp2_model(model) or keep_fused_experts in _process_quantized_modules is consistent with that.
• The _to_normal memory finding from the third pass is fixed: the DTensor branch is now guarded by _loc.is_inference(), so a no_grad caller (hf_ptq) copies nothing and peak is no longer the shard twice. The from_local(shape=…, stride=…) rewrap is still right for uneven and zero-size shards.
• Undeclared ties are a deliberate difference, not a gap. _write_even_shards clones rather than drops an aliased name ("every name has to reach the index"), so unlike save_pretrained/the streaming path this writer keeps both names. I went looking for a safetensors shared-memory failure here; the ptrs set at distributed_writer.py:563 already covers it, and declared ties are dropped by TiedWeightMap consistently on every rank.
• The input-scale reduce runs on _input_scale values that are plain tensors (registered after fully_shard, so FSDP2 never wraps them), which is what makes dist.all_reduce on the stack correct; amax/448 is monotonic, so reducing the post-fold scale equals reducing the amax, and no weight is quantized against a scale the checkpoint does not carry.
• broadcast_from is short-circuited by is_fsdp2_sharded, which comes from the recorded plan and is therefore identical on every rank; the save_modelopt_state refusal is raised before pack_quantized_weights mutates anything and before the try, so a refusal leaves the model untouched and no rank stranded in a collective; the extra_state_dict collision check runs inside all_reduce so rank 0’s finding raises everywhere.
• _write_base_config matches save_non_weight_artifacts on config dtype (both write model.config as-is), so the no-gather config.json does not diverge from the streaming paths’. File-write ordering is safe: copy_off_index_safetensors skips in-index names and _copy_non_model_files_from_source skips files the export wrote, so neither can clobber the shards or the embedded quantization_config.
• _writes_extra, rank and the _copy_non_model_files_from_source tail condition all follow the plan consistently; a replicated multi-rank export keeps its pre-existing "every rank writes everything" behaviour.
• Tests: the two-rank job’s ordering is deliberate and correct (the rank-0 byte comparison is last, after every collective, so a failing assert strands nobody); the mixed-precision KV case (_FP8_MIXED_KV_CFG) covers the _get_kv_cache_postprocess_config route end to end; the GPU test reads safetensors headers directly, halves the packed last dim for U8, and asserts every routed expert is quantized rather than just one. Dropping test_fused_experts_fold_in_a_weight_window… in favour of the Qwen2-MoE CPU case trades a monkeypatched call-sequence assertion for the observable effect — a net gain.
• PR size within budget: 355 insertions outside tests/ and docs/. CHANGELOG entries are user-facing, one or two sentences, and filed under the right sub-sections.

Assessment — risk: low-moderate. The writer stays opt-in and the default export is untouched, the plan-gate dependency that dominated the last pass is now satisfied inside this slice, and the structural failure modes (dropped key, mis-assembled shard, divergent ranks, duplicate name, missing index) are each asserted or explicitly refused. What is left is one behavioural delta against the path this writer substitutes for — a supported option becoming an exception, with a comment that states the opposite of what the sibling path does — and one knob-hygiene suggestion. Neither threatens checkpoint correctness; the first is worth closing because the env var is advertised as a drop-in.

🤖 Generated with Claude Code

Claude review — [16/21] No-gather sharded HF checkpoint write (re-review after the 2026-10-27 re-cut)

Scope: full review (the trigger comment carried no scoping instructions). 8 changed files, +1076/−117. Reviewed modelopt/torch/export/distributed_writer.py and modelopt/torch/export/unified_export_hf.py in full, plus examples/hf_ptq/README.md, CHANGELOG.rst, and all four test files (test_distributed_writer.py, test_unified_export_hf.py, test_fused_experts_keep_fused.py, and the new tests/gpu/torch/export/test_distributed_hf_export.py). Pulled in for context: _materialize_owned / _merge_groups / _revert_whole_sd / _write_even_shards from the base branch, moe_utils._split_local_fused_module, hf_export_handlers._export_weight / _export_fused_experts_module / _hold_fsdp2_weight_window, quant_utils.postprocess_state_dict / _get_kv_cache_postprocess_config, utils/distributed.DistributionPlan / broadcast_from / all_reduce, save_non_weight_artifacts, _write_hf_export_config and copy_non_model_files.

Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2 (both new)

1. unified_export_hf.py:1909-1928 — _write_base_config is a near-copy of the existing save_non_weight_artifacts; reusing it would also let the widened guard at L2182 revert to its original form.
2. distributed_writer.py:765-767 — reshape(()) on the expert input scales sits outside any all_reduce wrapper, so a rank-dependent scale shape or device raises on one rank while its peers wait in dist.all_reduce. Folding the shapes into the key agreement you just added closes it.

Both are non-blocking.

Prior review (2026-10-25) — status of its three suggestions

• Input-scale reduce ordering — fixed. The new all_reduce("expert input-scale keys", ...) makes every rank agree on the key list before the reduce, and a difference now fails on every rank naming the keys. The remaining gap is finding 2 above, which is the narrower residue of the same concern.
• EP scaffolding (ep_rank / dp_idx / total_experts / write_experts) — unchanged, and I am withdrawing the dead-code reading: the next PR in the series (#2753) is the expert-parallel export that sets ep_rank, so this is the seam it lands on rather than leftover. The one thing still worth a one-line fix is the comment "dp_idx is the EDP-group rank resolved above", which describes a resolution that does not happen on this path.
• get_rank() without is_initialized() (L1981, L2026) — unchanged. Only reachable by recording dense="fsdp" in a single process, so I am not re-raising it.

What I checked this round

• The two new behaviour fixes. Rank 0's modelopt_state save mirrors the FSDP2 streaming branch exactly (same guard, same filename, same rank), so the two paths cannot drift. _merge_groups(model, local_sd) is passed the full dict, experts included, but _materialize_owned filters each group through unit_of, which holds dense keys only; expert-only groups collapse to () and never enter unit_bytes, and a mixed group reduces to its dense members — identical on every rank, since the dense key set is proven identical by the gathered metadata. An expert-side merge (experts.<i>.gate_proj + up_proj) is already co-located by construction, so nothing is lost by the filter.
• Packing a sharded model. The relaxed _process_quantized_modules assertion opens the door for dense modules too, not only fused experts — but _export_weight opens an fsdp2_aware_weight_update window per FSDP2 unit whenever keep_fused_experts and plan.is_fsdp_dense, so dense weights are whole when packed. That is what keeps NVFP4's weight_scale_2 (a weight-amax reduction) from being computed over a local shard and silently mis-scaling the concatenated tensor. Experts stay Shard(0) and are folded in the unit's window, so FSDP2's next reshard cannot undo the fold.
• Collective symmetry end to end. use_no_gather_writer comes from broadcast_from on rank 0, so the two writers' different collective sequences cannot be mixed; the two-rank test sets the knob on rank 0 only. Within the writer, the key agreement, _materialize_owned's single metadata gather, _revert_whole_sd, the extra_state_dict overlap check and _write_even_shards's closing all_reduce are all unconditional and all route single-rank failures through modelopt.torch.utils.distributed.all_reduce. A rank-0-only failure in _write_base_config or _write_hf_export_config still reaches the finally: barrier(), so it raises rather than hangs.
• Checkpoint completeness vs save_pretrained. config.json + generation_config.json from _write_base_config, the index from _write_even_shards, hf_quant_config.json and the embedded quantization_config from _write_hf_export_config, tokenizer/processor/remote code/chat templates from _copy_non_model_files_from_source. The guard not (_offloaded or (is_fsdp2_sharded and not use_no_gather_writer)) evaluates correctly for all four path combinations, and the copy runs after the config writes, which copy_non_model_files documents as required. The ordering difference for _sanitize_generation_config_for_save (after config.save_pretrained here, before it in save_non_weight_artifacts) is harmless — it touches model.generation_config only.
• Quant-config / weight-name agreement. distributed_save_hf_checkpoint returning reverted, with _revert_quant_config_names_best_effort gated on it, is the tighter of the two conventions in this file (the streaming branch calls it unconditionally), and it is what keeps exclude_modules naming modules the way the weights do when the reverse is all-or-nothing refused.
• KV cache and dtypes. _get_kv_cache_postprocess_config is the same helper the other three HF exporters use, so MIXED_PRECISION resolves to the per-layer map here too; maxbound=448 matches the sibling call sites. The reduced scalars land as 0-dim float32, which is what the export handlers register.
• Backward compatibility. The plan requirement is the one break, and it is in CHANGELOG.rst under Backward Breaking Changes with the set_distribution_plan migration named. I traced every export_hf_checkpoint caller: hf_ptq.py records the plan right after parallel_load_and_prepare_fsdp2, the diffusers callers pass a pipeline (not an nn.Module, so the plan read is skipped), and the llm_sparsity callers are single-process. The new keep_fused_experts keyword defaults to the old behaviour on both pack_quantized_weights and _process_quantized_modules.
• Tests. The new unit cases run under the session-scoped cpu_dist_workers fixture, which CLAUDE.md exempts from the no-subprocess rule; CodeRabbit's two earlier notes are both addressed (the rank-0 byte-identical assertion is now explicitly last, "after every collective", and the test import moved to module scope). The divergence cases run under bounded_collectives(), so a regression fails fast instead of hanging CI. The GPU test reads safetensors headers directly rather than loading, skips cleanly under 2 GPUs, and its tie_word_embeddings case catches a doubly-written lm_head.weight through the unexpected assertion. The deleted weight-window test is justified and subsumed by the Qwen2-MoE end-to-end case.
• PR hygiene. 359 insertions under modelopt/ + examples/, inside the ~500-line budget. No prompt-injection-style text anywhere in the diff.

Risk: low. The new writer is opt-in behind MODELOPT_FSDP2_EXPORT=nogather and the default FSDP2 export is untouched; the shared-code changes are additive with old defaults. Both this round's findings are diagnostics and simplicity, not behaviour, and the three fixes since the last review each closed a real asymmetry. Approving.

🤖 Generated with Claude Code